modelopt/torch/speculative/plugins/megatron_eagle.py (1)

531-551: EAGLE-3 first-layer specialization looks correct; please keep hook handles to allow cleanup.

The pre-hook and widened QKV projection match the 2h→(q+2kv) requirement. Store the returned handles so they can be removed during teardown or reconfiguration to avoid duplicate hooks in long-lived processes.

- last_layer.register_forward_hook(self._eagle3_layer_forward_hook)
+ self._hooks = getattr(self, "_hooks", [])
+ self._hooks.append(last_layer.register_forward_hook(self._eagle3_layer_forward_hook))

- self_attention.register_forward_pre_hook(self._eagle3_attention_forward_pre_hook)
+ self._hooks.append(
+     self_attention.register_forward_pre_hook(self._eagle3_attention_forward_pre_hook)
+ )

Optionally add a small def remove_hooks(self): that removes and clears self._hooks.

tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py (3)

73-84: Ensure convert() receives the expected structure and prevent accidental in-place edits.

If EAGLE*_DEFAULT_CFG is a list like [("eagle", {"config": {...}})], mutate the nested dict only after deepcopy and keep the outer structure intact. Otherwise, wrap into [("eagle", cfg)]. Please verify the expected schema for mtsp.convert.

-    elif algo in {"eagle1", "eagle3"}:
-        mtsp_config = ALGO_TO_CONFIG[algo]
+    elif algo in {"eagle1", "eagle3"}:
+        mtsp_config = deepcopy(ALGO_TO_CONFIG[algo])
@@
-        model = mtsp.convert(model, mtsp_config)
+        # If mtsp_config is a bare dict (not a list of (name, cfg)), wrap it:
+        model = mtsp.convert(model, mtsp_config if isinstance(mtsp_config, list) else [("eagle", mtsp_config)])

---

88-89: Minor: fix error message formatting.

Use an f-string or a raw literal without braces to avoid confusion.

-        raise ValueError("Only algo={eagle1, eagle3, medusa} are supported!")
+        raise ValueError("Only algo in {eagle1, eagle3, medusa} are supported!")

---

145-154: Param set expansion is good; please update the skip guard accordingly.

Below, the skip guard still checks algo == "eagle", which won’t trigger for eagle1/eagle3. Update to cover both variants.

-    if algo == "eagle":
+    if algo.startswith("eagle"):
         try:
             import megatron.core.post_training  # noqa: F401
         except ImportError:
             pytest.skip("megatron.core.post_training not found")

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

modelopt/torch/speculative/plugins/megatron_eagle.py (2)

527-528: Hook captures layer output, not truly “pre-norm”; store handle to avoid duplicate registrations

• The comment says “pre-norm hidden_state,” but a forward hook on the last TransformerLayer returns that layer’s output (pre-final-layernorm of the block, which is fine if that’s the intent). If you truly want the module input (pre-layernorm of the last layer), switch to using the hook’s input[0] inside _eagle3_layer_forward_hook. Otherwise, please update the comment to “pre-final-layernorm” to avoid confusion.
• Minor: Keep the returned hook handle to prevent accidental duplicate registrations (e.g., if this ctor is ever run twice on the same instance pattern).

Proposed small refactor to retain the hook handle:

- last_layer = self.decoder.layers[-1]
- last_layer.register_forward_hook(self._eagle3_layer_forward_hook)
+ last_layer = self.decoder.layers[-1]
+ self._last_layer_hook = last_layer.register_forward_hook(self._eagle3_layer_forward_hook)

If you do want “pre-layernorm of the last layer” (not just pre-final-layernorm), adjust the hook implementation (outside this hunk) like:

def _eagle3_layer_forward_hook(self, module, input, output):
    # input[0] is the hidden_states entering the layer (pre-layernorm in Pre-LN stacks)
    hidden_states = input[0].detach().clone()
    self._next_hidden_states_input = hidden_states

---

530-531: Guard first-layer access and explain specialization intent

Accessing self.decoder.layers[0] assumes num_layers >= 1. Add an explicit check and a short comment to clarify that this specialization is only for EAGLE-3’s first decoder layer.

- # The first EAGLE3 layer needs to be specialized.
- layer = self.decoder.layers[0]
+ # The first EAGLE3 layer needs to be specialized (EAGLE-3 requires it).
+ assert len(self.decoder.layers) >= 1, "EAGLE-3 requires at least 1 decoder layer"
+ layer = self.decoder.layers[0]

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, both EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG in config.py already use deepcopied configurations (either directly via deepcopy() or through variables that were already deepcopied), making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.359Z
Learning: In the TensorRT-Model-Optimizer codebase, the EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG configurations in config.py are already deepcopied, so additional deepcopy calls are not needed when using ALGO_TO_CONFIG mapping in tests.